import time
import json
from core.coze_bot import CozeBot as CB
from core.docx_converter import DocxConverter as DC

"""
文章生成模型，传入写作机器人模型cb、doc文档对象dc、文章标题、写作模型等信息，生成文章并保存到docx。
其中，写作模型，有model_one、model_two、model_three....等大概10个写作模型，每个写作模型逻辑均不同。请规划下代码，并优化下代码。
model_one:一篇文章有多个独立的部分，相互之间无关系。
model_two:一篇文章有多个连续的章节，相互之间有关系，但不需要xml转换，无子标题。在提示词方面，会有较大区别。
model_three:一篇文章有多个连续的章节，相互之间有关系，需要xml转换，有多级子标题。
"""


class ArticleModelApi:

	def __init__(self, title, cb=None, dc=None, model=None):
		self.title = title
		self.cb = cb  # 写作机器人实例
		self.dc = dc  # 文档实例
		self.model = model

	# 开始按照流程工作,
	def create_one_article(self):
		# print(self.model)
		if int(self.model['model_type']) == 1:
			self._model_one_write()
		elif int(self.model['model_type']) == 2:
			# 另一个逻辑
			self._model_two_write()
		elif int(self.model['model_type']) == 3:
			self._model_three_write()
	# 开始创作文字：model_one
	def _model_one_write(self):
		# ================分三步获取大纲===========================
		# 1、获取prompt
		prompt_requirements = self.model['prompt'][0]
		# 2、序列化,得到可用于请求的bot
		chat_history = [
			{"role": "user", "content_type": "text", "content": prompt_requirements['system_content']},
			{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！",
			 "content_type": "text"}
		]
		query = prompt_requirements['query'].format(title=self.title)
		title_format = {
			"chat_history": chat_history,
			"query": query
		}
		# 3、问机器人，返回title和requirements
		# print(title_format)
		ress = self.cb.ask_bot(title_format)
		# print(ress)
		if not ress:
			return False
		article_requirements = json.loads(ress)

		# print(f"问机器人，返回title和requirements{article_requirements}")
		if article_requirements:  # 循环写作
			self._model_one_write_loop(article_requirements)

	def _model_one_write_loop(self, article_requirements=[]):
		# print(f"文章写作要求")
		i = 1
		for requirement in article_requirements["requirements"]:
			error_num = 0
			while True:
				if error_num >= 5:
					break
				# ================分三步获取正文===========================
				# 1、获取prompt
				prompt_content = self.model['prompt'][1]
				# 2、序列化,得到可用于请求的bot
				chat_history = [
					{"role": "user", "content_type": "text", "content": prompt_content['system_content']},
					{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！",
					 "content_type": "text"}
				]
				# print(prompt_content)
				query = prompt_content['query'].format(title=self.title, requirement=requirement)
				content_format = {
					"chat_history": chat_history,
					"query": query
				}
				# print(content_format)
				# 3、问机器人,返回了文章内容
				article_content = self.cb.ask_bot(content_format)
				# print(article_content)
				if not article_content:
					error_num += 1
					continue

				if self.model['code'] == 'english_article':#不需要xml转换
					docx_res = self.dc.add_text_to_docx(text_string=article_content, section_title=str(i))
					if docx_res:
						print(f"已完成《{self.title}》的第{i}部分，并插入文档中.......")
						i += 1
						break
				# ================分三步转换成xml文档===========================
				# 1、获取prompt



				prompt_xml = self.model['prompt'][2]
				# 2、序列化,得到可用于请求的bot
				chat_history = [
					{"role": "user", "content_type": "text", "content": prompt_xml['system_content']},
					{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！",
					 "content_type": "text"}
				]
				query = prompt_xml['query'].format(content=article_content)
				xml_format = {
					"chat_history": chat_history,
					"query": query
				}
				# print(xml_format)
				# 3、问机器人,返回了文章内容
				xml_content = self.cb.ask_bot(xml_format)
				# print(xml_content)
				if not xml_content:
					error_num += 1
					continue
				# ================把xml文档插入到文档对象中去===========================

				docx_res = self.dc.add_xml_to_docx(xml_string=xml_content, section_title=str(i))
				if docx_res:
					print(f"已完成《{self.title}》的第{i}部分，并插入文档中.......")
					i += 1
					break
		self.dc.save_output()

	def _model_two_write(self):
		# ================分三步获取大纲===========================
		# 1、获取prompt
		prompt_requirements = self.model['prompt'][0]
		# 2、序列化,得到可用于请求的bot
		chat_history = [
			{"role": "user", "content_type": "text", "content": prompt_requirements['system_content']},
			{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！","content_type": "text"}
		]
		query = prompt_requirements['query'].format(title=self.title)
		title_format = {
			"chat_history": chat_history,
			"query": query
		}
		# 3、问机器人，返回title和requirements
		# print(title_format)
		res11 = self.cb.ask_bot(title_format)
		# print(res11)
		try:
			article_requirements = json.loads(res11)
		except Exception as e:
			time.sleep(10)
			print(res11,e)
			return self._model_two_write()
		# print(article_requirements)
		# print(f"问机器人，返回title和requirements{article_requirements}")
		if article_requirements:  # 循环写作
			# print(article_requirements)
			self._model_two_write_loop(article_requirements)

	def _model_two_write_loop(self, article_requirements=[]):
		# print(f"文章写作要求")
		i = 1
		for requirement in article_requirements["section_requirements"]:
			error_num = 0
			while True:
				if error_num >= 5:
					break
				# ================分三步获取正文===========================
				# 1、获取prompt
				prompt_content = self.model['prompt'][1]
				# 2、序列化,得到可用于请求的bot
				chat_history = [
					{"role": "user", "content_type": "text", "content": prompt_content['system_content']},
					{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！", "content_type": "text"}
				]
				# print(prompt_content)
				query = prompt_content['query'].format(title = self.title,overall_requirement=article_requirements["overall_requirements"], section_requirement=requirement)
				content_format = {
					"chat_history": chat_history,
					"query": query
				}
				# print(content_format)
				# 3、问机器人,返回了文章内容
				article_content = self.cb.ask_bot(content_format)
				# print(article_content)
				if not article_content:
					error_num += 1
					continue

				docx_res = self.dc.add_text_to_docx(text_string=article_content )
				if docx_res:
					print(f"已完成《{self.title}》的第{i}部分，并插入文档中.......")
					i += 1
					break
		self.dc.save_output()

	def _model_three_write(self):
		# ================分三步获取大纲===========================
		# 1、获取prompt
		prompt_requirements = self.model['prompt'][0]
		# 2、序列化,得到可用于请求的bot
		chat_history = [
			{"role": "user", "content_type": "text", "content": prompt_requirements['system_content']},
			{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！",
			 "content_type": "text"}
		]
		query = prompt_requirements['query'].format(title=self.title)
		title_format = {
			"chat_history": chat_history,
			"query": query
		}
		# 3、问机器人，
		res11 = self.cb.ask_bot(title_format)
		# print(res11)
		try:
			article_requirements = json.loads(res11)
		except Exception as e:
			print(res11, e)
			return self._model_three_write()
		if article_requirements:  # 循环写作
			self._model_three_write_loop(article_requirements)

	def _model_three_write_loop(self, article_requirements=[]):
		# print(f"文章写作要求")
		i = 1
		for requirement in article_requirements["section_instructions"]:
			error_num = 0
			while True:
				if error_num >= 10:
					break
				# ================分三步获取正文===========================
				# 1、获取prompt
				prompt_content = self.model['prompt'][1]
				# 2、序列化,得到可用于请求的bot
				chat_history = [
					{"role": "user", "content_type": "text", "content": prompt_content['system_content']},
					{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！",
					 "content_type": "text"}
				]
				query = prompt_content['query'].format(title=self.title,section_title=requirement['section_title'],overall_instructions=article_requirements["overall_instructions"],instructions=requirement['instructions'])
				content_format = {
					"chat_history": chat_history,
					"query": query
				}
				# print(content_format)
				# 3、问机器人,返回了文章内容
				article_content = self.cb.ask_bot(content_format)
				# print(article_content)
				if not article_content:
					error_num += 1
					continue
				# print(article_content)
				# return 1
				# ================分三步转换成xml文档===========================
				# 1、获取prompt
				prompt_xml =  self.model['prompt'][2]
				# 2、序列化,得到可用于请求的bot
				chat_history = [
					{"role": "user", "content_type": "text", "content": prompt_xml['system_content']},
					{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！", "content_type": "text"}
				]
				query = prompt_xml['query'].format(content=article_content)
				xml_format = {
					"chat_history": chat_history,
					"query": query
				}
				# 3、问机器人,返回了文章内容
				xml_content = self.cb.ask_bot(xml_format)
				print("-----------------------------------------------------------------------------------------------")
				# print(article_content)
				# print(xml_content)
				if not xml_content:
					error_num += 1
					continue
				# ================把xml文档插入到文档对象中去===========================

				docx_res = self.dc.add_xml_to_docx(xml_string=xml_content)
				if docx_res:
					print(f"已完成《{self.title}》的第{i}部分，并插入文档中！")
					i += 1
					break
		self.dc.save_output()


if __name__ == "__main__":
	cb = CB()
	dc = DC(title="商业模式创新专家年度总结报告")
	model = {
		"id": 5,
		"code": "gzzj",
		"name": "工作总结（测试中）",
		"remark": "先大纲，分部写，XML转换",
		"model_type": 3,
		"model_conf": "{}",
		"create_time": "2024-08-08 21:00:44",
		"update_time": '',
		"prompt": [
			{
				"id": 13,
				"model_id": 5,
				"code": "1",
				"remark": "工作总结-标题分析",
				"system_content": "# 角色\r\n你是一位精通且卓越的公文写作策划大师，能够精准应对各种复杂的写作需求，全面且细致地开展公文写作策划工作，并以规范的 json 格式进行输出。\r\n\r\n## 技能\r\n### 技能 1: 深入挖掘写作关键要素\r\n1. 全面且深入地探究写作主题及目标，精准界定公文的性质、受众类别以及期望达成的具体成效。\r\n2. 充分考虑公文所属领域的独特特点和相关的规范标准要求。\r\n\r\n### 技能 2: 精确设定总篇幅\r\n1. 综合权衡公文内容的丰富程度与表达需要，设定不低于 3000 字的总字数。\r\n2. 科学合理地规划各章节的字数分配比例，突出重点部分，做到详略安排得当。\r\n\r\n### 技能 3: 精心划分章节结构\r\n1. 依据公文的严密逻辑架构和核心内容要点，进行科学合理的章节划分。\r\n2. 控制每个章节的字数在 1000 字以内，保证结构紧凑、条理清晰。\r\n3. 章节数量控制在 3 至 8 个之间。\r\n\r\n### 技能 4: 明确各章节具体要求\r\n1. 针对每个章节，清晰明确主题、核心观点以及重点阐述的内容。\r\n2. 严格规定章节内的论证方式、引用资料的规范和语言风格。\r\n3. 适时适量地插入具有代表性的实例。\r\n\r\n\r\n## 输出示例\r\n\"\"\"\r\n{\r\n\"title\":\"年度销售工作总结报告\",\r\n\"word_number\":\"4000\",\r\n\"overall_instructions\":\"整体要求是：撰写一份高质量的、深度的、惊艳的、感情丰富的总结报告，全面回顾本年度的销售工作，报告分为六大章节，分别是开头、总结前阶段成果、讲述完成任务采取的措施、存在的不足与问题、下一步工作计划、结语，采用1, 1.1, 1.1.1的层次化编号体系，确保内容结构清晰，每章标题精炼、对仗工整且风格统一。整体风格应追求高质量、引人入胜。\",\r\n\"section_instructions\":[{\r\n\"section_number\":\"1\",\r\n\"section_title\":\"1.开篇概览，启航新程\",\r\n\"instructions\":\"第一部分为开头，开篇需精炼概述，概括性总结过去一年销售工作的辉煌成就与背景环境，营造总结氛围，篇幅约300字，一个段落，不需要标题，如：'岁末之际，回望征途，销售军团在公司领航者的睿智部署与全体将士的并肩奋斗下，铸就了业绩的辉煌篇章。此刻，我们提笔总结，既为沉淀过往，更为照亮未来之路。'\"\r\n},{\r\n\"section_number\":\"2\",\r\n\"section_title\":\"2.辉煌战绩，数据见证\",\r\n\"instructions\":\"总结前阶段成果，详细列举并量化销售业绩，包括但不限于利税贡献、生产总值、客户签约量、产品开发成果等，通过多级标题深入阐述，确保数据详实，内容丰富，总字数超过1000字，要有子标题，如'1.1 利税贡献显著'等。\"\r\n},{\r\n\"section_number\":\"3\",\r\n\"section_title\":\"3.策略与智慧，铸就胜利的篇章\",\r\n\"instructions\":\"讲述完成任务采取的措施。深入分析实施成功的策略与措施，涵盖企业管理创新、团队建设优化、品牌宣传力度、市场营销策略等方面，强调事实依据与数据支撑，分为两级标题，总字数不少于1000字，要有子标题，如'2.1 健全机构，强化组织领导；2.2 创新方法，拓宽盈利渠道'等。\"\r\n},{\r\n\"section_number\":\"4\",\r\n\"section_title\":\"4.自省与超越，攀登更高的山峰\",\r\n\"instructions\":\"存在的不足与问题。深刻剖析工作中存在的问题与不足，从思想认知、学习成长、管理体系、执行方法等多维度探讨原因，至少提出三个方面的不足，并给出初步反思，总字数不少于1000字，要有子标题，如'3.1 认知深度待提升'等。\"\r\n},{\r\n\"section_number\":\"5\",\r\n\"section_title\":\"5.展望未来，绘制梦想的蓝图\",\r\n\"instructions\":\"下一步工作计划。结合公司总体战略与自身工作实际，明确未来一年的销售目标与重点任务，提出具体可行的工作计划与措施，总字数不少于400字，要有子标题，如'4.1 品质为先定根基；4.2 渠道为王固市场'等。\"\r\n},{\r\n\"section_number\":\"6\",\r\n\"section_title\":\"6.携手并进筑辉煌\",\r\n\"instructions\":\"结尾。以饱满的热情展望新的一年，表达团队决心与愿景，如：'步入新的一年，我们将紧密围绕公司战略核心，秉承拼搏精神，以更加坚定的步伐，攻克市场难关，开创销售新篇章，为实现公司宏伟目标贡献力量！' ，展现积极向上的态度与坚定的信念，不需要子标题。不少于300字。\"\r\n}]}\r\n\"\"\"\r\n## 限制:\r\n- 仅专注于公文写作策划相关事宜，不涉足其他无关领域。\r\n- 务必严格遵循上述设定的格式和要求进行策划，严禁随意偏离。\r\n- 策划内容必须具备合理性、可行性和高度的专业性。",
				"query": "请分析这个写作任务：《{title}》。\r\n按照输出格式示例，直接输出json。",
				"create_time": "2024-07-15 22:36:31",
				"update_time": 'null',
				"step_sort": 1
			},
			{
				"id": 14,
				"model_id": 5,
				"code": "1",
				"remark": "工作总结-内容生产",
				"system_content": "# 角色\r\n你是一位卓越的公文写作专家，能够依据用户提供的总体需求和具体章节要求，迅速且高质量地创作出情感真挚、条理明晰、深度足够且完全符合公文规范的公文章节内容。\r\n\r\n## 技能\r\n### 技能 1: 精准把握需求\r\n1. 全面深入地解读用户给出的总体及章节特定要求，确保准确无误、清晰透彻。\r\n2. 针对不明确之处，进行合理且恰当的推测与补充，无需向用户求证。\r\n\r\n### 技能 2: 精心构建章节\r\n1. 严格依照公文的标准格式与规范进行创作。\r\n2. 设计引人入胜的章节标题，确保结构清晰有序。\r\n3. 采用丰富、精准且恰当的词汇及表达方式，让内容富有情感与深度。\r\n4. 保障章节内容层次分明，逻辑严密无漏洞。\r\n5. 完成创作后，认真核查语言表述和格式方面的错误。\r\n\r\n## 限制:\r\n- 仅专注于公文写作相关事务，绝不参与与公文写作无关的活动。\r\n- 输出内容务必严格遵循公文的格式与规范，不容有任何偏差。\r\n- 所创作的章节内容必须紧密围绕给定要求，严禁随意拓展。",
				"query": "现在让我们来协同编写一份名为《{title}》的高质量的总结报告。其整体要求如下：\r\n\"\"\"\r\n{overall_instructions}\r\n\"\"\"\r\n接下来，是展现您才华与创意的时刻，请您精心撰写其中的章节《{section_title}》，此章节的具体指导方针与期待成果如下：\r\n\"\"\"\r\n{instructions}\r\n\"\"\"\r\n在此章节中，我们尤为注重内容的深度、创新性与独特性，期许您的文字能够独树一帜，不仅准确传达信息，更能激发思考，令人耳目一新。请确保文字质量上乘，既符合学术或行业规范，又不失灵动与韵味。同时，请确保篇幅充实，以详尽而精炼的论述支撑您的观点与发现。期待你惊艳的表现，让我们共同见证这份总结报告的精彩诞生。\r\n请直接输出章节内容，包括章节标题。无需附加其他内容。\r\n",
				"create_time": "2024-07-15 22:39:30",
				"update_time": 'null',
				"step_sort": 2
			},
			{
				"id": 15,
				"model_id": 5,
				"code": "1",
				"remark": "工作总结-格式转换",
				"system_content": "# 角色\r\n您是一位专业且高效的文章至XML格式转换器，擅长将繁杂的文章内容转化为结构清晰、层次分明的XML文档。\r\n\r\n## 技能\r\n### 技能 1：精准转换文章为 XML 格式\r\n1. 精准识别文章中的各级标题编号（如1、2、3为一级，1.1、1.2为二级，依此类推），并准确无误地将它们映射至对应的<title1>、<title2>等标签内，确保文档结构的严谨性。\r\n2. 自动检测并规范地将文章内容中的每一个段落包裹在独立的<paragraph>标签中，提升内容的可读性与可管理性。\r\n3. 确保各级标题与段落之间的嵌套关系准确无误，完美还原文章原有的逻辑结构与层次感。\r\n\r\n## 输出示例(title1、title2、title3、paragraph之间无包含关系)：\r\n<content>  \r\n    <title1>1. 一级标题 1</title1>  \r\n    <title2>1.1 二级标题</title2>  \r\n    <paragraph>这是一段内容</paragraph>\r\n    <title1>2. 一级标题 2</title1>\r\n    <paragraph>正文部分</paragraph>\r\n</content>\r\n\r\n## 约束\r\n- 禁止添加任何文字，只可适当删减。\r\n- 对文章标题的识别采取保守方式，若无法明确判定为标题，就按段落处理。\r\n- 仅交流与文章内容转换为 XML 格式有关的话题。\r\n- 尽最大努力精准依照文章的层次结构生成标题标签与段落的嵌套关系。\r\n- 务必保证所有的标签都正确闭合，避免引发解析错误。",
				"query": "请将以下文章内容准确无误地转换为XML格式，过程中切勿添加除必要的XML标签外的任何额外内容。同时，请尽可能细致地区分段落，确保内容的结构化呈现。\r\n\r\n在转换过程中，请注意以下几点关键规则：\r\n\r\ntitle1、title2、title3 以及 paragraph 标签之间不存在直接的包含关系，它们各自独立，用于标记不同级别的标题与正文段落。\r\n转换完成后，务必仔细检查XML结构的完整性，确保所有开启的标签都有对应的闭合标签，避免解析错误。\r\n所有的标题（title1、title2、title3）和段落（paragraph）末尾，若原文中未包含标点符号，请依据上下文适当添加，以保证内容的准确性和可读性。\r\n现提供文章内容如下：\r\n\"\"\"\r\n{content}\r\n\"\"\"\r\n请基于上述规则与提示，生成并返回完整的XML结构。",
				"create_time": "2024-07-15 23:12:59",
				"update_time": 'null',
				"step_sort": 3
			}
		]
	}
	AM = ArticleModelApi("商业模式创新专家年度总结报告", cb, dc, model)
	AM.create_one_article()
